#!/usr/local/bin/python

import sys
import psycopg2
import os

if len(sys.argv) != 1:
    sys.exit("""
    Print on stdout the table schemas for the tables found in
    sblab.main.filing_order.table_name.
    
    Edit filing_order to include/exclude tables or change the order they are
    printed out.
    
    USAGE
        get_pgschema.py
    
    REQUIRES
        - File ~/.psycopgpass with the log in details to pass to psycopg2.connect()
             """)

# -----------------------------------------------------------------------------
def get_psycopgpass():
    """
    Read file ~/.psycopgpass to get the connection string to pass to
    psycopg2.connect()
    Example usage:
        conn= psycopg2.connect(get_psycopgpass())
        cur= conn.cursor()
    """
    conn= open(os.path.join(os.getenv("HOME"), '.psycopgpass'))
    conn= conn.readlines()
    conn_args= [x.strip() for x in conn if x.strip() != '' and x.strip().startswith('#') is False][0]
    return(conn_args)
# ------------------------------------------------------------------------------

conn= psycopg2.connect(get_psycopgpass())
cur= conn.cursor()

sql= """ SELECT update_filing_order();
    CREATE TEMP TABLE tconstraints AS(
    -- See http://stackoverflow.com/questions/1152260/postgres-sql-to-list-table-foreign-keys for where this query was taken from
    SELECT DISTINCT 
        tc.constraint_name, 
        tc.table_name, 
        kcu.column_name, 
        tc.constraint_type,
        ccu.table_name AS foreign_table_name,
        ccu.column_name AS foreign_column_name 
    FROM 
        information_schema.table_constraints AS tc 
        JOIN information_schema.key_column_usage AS kcu ON tc.constraint_name = kcu.constraint_name
        JOIN information_schema.constraint_column_usage AS ccu ON ccu.constraint_name = tc.constraint_name
    WHERE tc.table_name in (select table_name from filing_order where filing_order > 0 and filing_order is not null))"""

cur.execute(sql)

sql= """SELECT DISTINCT c.table_name,
           c.column_name,
           c.data_type,
           c.ordinal_position,
           filing_order.filing_order,
           sort_distinct(array_agg(tconstraints.constraint_type ||  
                     CASE WHEN tconstraints.constraint_type != 'FOREIGN KEY' THEN '' ELSE ' ' || tconstraints.foreign_table_name END || 
                     CASE WHEN tconstraints.constraint_type != 'FOREIGN KEY' THEN '' ELSE '.' || tconstraints.foreign_column_name END)) AS constraint,
           c.is_nullable
    FROM information_schema.columns AS c
        JOIN filing_order on filing_order.table_name = c.table_name
        LEFT JOIN tconstraints ON c.table_name = tconstraints.table_name AND c.column_name = tconstraints.column_name 
    WHERE filing_order > 0 and filing_order is not null
    GROUP BY c.table_name, c.column_name, c.ordinal_position, filing_order.filing_order, c.data_type, c.is_nullable
    ORDER BY filing_order DESC, ordinal_position;
    """
cur.execute(sql)
pglist= cur.fetchall()

def getMaxLength(llist, indx):
    """ Get the max column length for a table
    llist: list of lists from which to group-by and extract the longest string
    indx:  column index for which to get the length
    
    Returns a dict of group names (keys) and max string length
    """
    collength= {}
    for line in llist:
        if not collength.has_key(line[0]):
            collength[line[0]]= len(line[indx])
        elif collength[line[0]] < len(line[indx]):
            collength[line[0]]= len(line[indx])
    return(collength)

def formatConstraint(xlist, sep= ' '):
    """ Format constraint xlist (list) for pretty printing.
    Join in a single string using sep as joiner
    """
    xcon= []
    for x in xlist:
        if x is None or x == 'PRIMARY KEY' or x == 'UNIQUE':
            xcon.append('')
        else:
            x= x.replace('FOREIGN KEY', '-->')
            xcon.append(x)
#    if 'UNIQUE' in xcon:
#        i= xcon.index('UNIQUE')
#        xcon.insert(0, xcon[i])
#        del(xcon[i+1])
#    if 'PRIMARY KEY' in xcon:
#        i= xcon.index('PRIMARY KEY')
#        xcon.insert(0, xcon[i])
#        del(xcon[i+1])
    xcon= [x for x in xcon if x != '']
    return(sep.join(xcon))

pglist2= []
for line in pglist:
    line= list(line)
    pglist2.append(line)
    if line[6] == 'YES':
        line[2]= line[2] + ' not null'
    if 'PRIMARY KEY' in line[5]:
        line[1]= '** ' + line[1]
    elif 'UNIQUE' in line[5]:
        line[1]= '++ ' + line[1]
    else:
        line[1]= '   ' + line[1] + '' 
    newk= []
    for x in line[5]:
        if x is None:
            pass
        elif x == 'PRIMARY KEY' or x == 'UNIQUE':
            pass
        elif x.startswith('FOREIGN KEY'):
            x= x.replace('FOREIGN KEY', '-->')
            newk.append(x)
        else:
            newk.append(x)
    line[5]= sorted(list(set(newk)))
    
colOffset= 0
for line in pglist2:
    if len(line[0]) > colOffset:
        colOffset= len(line[0])
colOffset= 2
dtypeOffset= colOffset + max([len(x) for x in [x[1] for x in pglist2]])+2 ## Column where data types should start
fkeyOffset= dtypeOffset + max([len(x) for x in [x[2] for x in pglist2]])+2 ## Column where keys should start

schemalist= []
curtable= ''
for line in pglist2:
    row= [' ']*200 ## Initialize an empty row
    if line[0] != curtable:
        curtable= line[0]
        schemalist.append(''.join([' ']*200))
        schemalist.append(''.join(['-']*80 + [' ']*120))
        schemalist.append(''.join(list(line[0]) + [' ']*(200 - len(line[0]))))
    row[colOffset: (colOffset + len(line[1]))]= list(line[1])
    row[dtypeOffset: (dtypeOffset + len(line[2]))]= list(line[2])
    if line[5] == []:
        schemalist.append(''.join(row))
    else:
        for i in range(0, len(line[5])):
            fkey= line[5][i]
            row[fkeyOffset: (fkeyOffset + len(fkey))]= list(fkey)
            if i == 0:
                schemalist.append(''.join(row))
            else:
                row[0:fkeyOffset]= ' '*fkeyOffset
                schemalist.append(''.join(row))
## Get coordinates of foreign keys:
fkeys= {} ## Each entry of this will be: {'line no. of fkey (int)': ['referenced column', 'line no. of referenced col.']}
for i in range(0, len(schemalist)):
    line= schemalist[i]
    if '-->' in line:
        fkey= line[(line.find('-->')+4):].strip()
        fkeys[i]= [fkey]
## Coordinates of associated primary keys
for k in fkeys:
    fkey= fkeys[k][0]
    (table, column)= fkey.split('.')
    ok=False
    n= 0
    for line in schemalist:
        if line.strip() == table:
            ok= True
        if ok and line.strip(' *+').startswith(column + ' '):
            fkeys[k]= fkeys[k] + [n]
        n += 1
#print(fkeys)
nfkey= len(fkeys)
keys= sorted(fkeys.keys())
drawOffset= 0
for line in schemalist:
    line= line.rstrip()
    if len(line) > drawOffset:
        drawOffset= len(line)

schemalist= [list(x) for x in schemalist]
prevend= 0
for s in keys:
    try:
        e= fkeys[s][1]
    except IndexError:
        sys.exit('ERROR:\nIt appears that index %(t)s has no index associated. Make sure that table holding %(t)s is present in table filing_order with a filing_order > 0 and filing_order not null' %{'t': fkeys[s][0]})
    if s > e:
        start= e
        end= s
        c= '|'
    else:
        start= s
        end= e
        c= '|'
#    if prevend >= start:
#        drawOffset += 1
#    prevend= end
    curOffset= drawOffset
    while True:
        trace= []
        for i in range(start-2, end+2):
            trace.append(schemalist[i][curOffset])
        if set(trace) == set([' ']):
            break
        else:
            curOffset += 2
    for i in range(start, end):
            schemalist[i][curOffset]= '|'
    schemalist[e][curOffset]= '<'
    schemalist[s][curOffset]= '+'
    oldend= end
for line in schemalist:
    print(''.join(line).rstrip())
print('')
print('                                                 +------- Key to symbols -------+')
print('                                                 |                              |')
print('                                                 | ** : Primary key(s)          |')
print('                                                 | ++ : Unique constraint       |')
print('                                                 | -->: Foreing key             |')
print('                                                 | <-+: Foreing key connections |')
print('                                                 +------------------------------+\n')

sql= "DROP TABLE tconstraints"
cur.execute(sql)
sys.exit()